fix(core): validate replay operation identity - #698

"""Tests for map summarized replay across a suspension — parity with JS PR #894.

Matrix (mirrors packages/aws-durable-execution-sdk-js-examples map/flat-summarized-replay):
- FLAT nesting, >256KB result  -> ReplayChildren rebuild across a suspension (the bug path)
- NESTED nesting, >256KB result -> control: context checkpoints exist, rebuild trivially correct
- FLAT nesting, small result    -> control: stays within one checkpoint, no rebuild at all
"""

import pytest
from aws_durable_execution_sdk_python.config import MapConfig, NestingType
from aws_durable_execution_sdk_python.context import DurableContext
from aws_durable_execution_sdk_python.execution import (
    InvocationStatus,
    durable_execution,
)
from aws_durable_execution_sdk_python.waits import Duration

from test.conftest import deserialize_operation_payload

ITEM_COUNT = 8
LARGE_ITEM_PAYLOAD_BYTES = 40 * 1024  # 8 x 40KB ≈ 320KB > 256KB limit
SMALL_ITEM_PAYLOAD_BYTES = 100  # 8 x 100B — far below the limit


def make_handler(nesting: NestingType, payload_bytes: int):
    @durable_execution
    def handler(_event, context: DurableContext):
        items = list(range(ITEM_COUNT))

        batch = context.map(
            inputs=items,
            func=lambda ctx, item, index, _: ctx.step(
                lambda _s: {"item": item, "payload": "x" * payload_bytes},
                name=f"resolve-{index}",
            ),
            name="resolve-pages",
            config=MapConfig(nesting_type=nesting),
        )

        live_result_count = context.step(
            lambda _: len(batch.get_results()), name="record-live-count"
        )

        context.wait(Duration.from_seconds(86400), name="suspend")

        replayed_result_count = context.step(
            lambda _: len(batch.get_results()), name="record-replayed-count"
        )
        replayed_items = sorted(r["item"] for r in batch.get_results())

        return {
            "live_result_count": live_result_count,
            "replayed_result_count": replayed_result_count,
            "replayed_items": replayed_items,
        }

    return handler


HANDLERS = {
    "flat-large": make_handler(NestingType.FLAT, LARGE_ITEM_PAYLOAD_BYTES),
    "nested-large": make_handler(NestingType.NESTED, LARGE_ITEM_PAYLOAD_BYTES),
    "flat-small": make_handler(NestingType.FLAT, SMALL_ITEM_PAYLOAD_BYTES),
}


@pytest.mark.example
@pytest.mark.parametrize("case", ["flat-large", "nested-large", "flat-small"])
def test_map_summarized_replay_matrix(case):
    """The batch reconstructed after suspend/resume must equal the live batch."""
    from aws_durable_execution_sdk_python_testing.runner import (
        DurableFunctionTestRunner,
    )

    with DurableFunctionTestRunner(handler=HANDLERS[case]) as runner:
        result = runner.run(input='"test"', execution_timeout=30)

    assert result.status is InvocationStatus.SUCCEEDED, (
        f"[{case}] expected SUCCEEDED, got {result.status}: {result.error}"
    )
    output = deserialize_operation_payload(result.result)
    assert output["live_result_count"] == ITEM_COUNT, f"[{case}] live count wrong"
    assert output["replayed_result_count"] == ITEM_COUNT, (
        f"[{case}] replayed count {output['replayed_result_count']} != {ITEM_COUNT} — "
        "summarized-replay rebuild lost terminal items"
    )
    assert output["replayed_items"] == list(range(ITEM_COUNT)), (
        f"[{case}] replayed items {output['replayed_items']}"
    )

"""Non-deterministic replay fails closed with diagnostic — parity with JS PR #867.

Two cases, mirroring replay-validation.composed.test.ts:
1. Divergent replay (STEP checkpointed, CHAINED_INVOKE replayed at the same
   position) -> FAILED, error type preserved as NonDeterministicExecutionError,
   never PENDING.
2. Stable control (identical replay) -> SUCCEEDED.
"""

import pytest
from aws_durable_execution_sdk_python.context import DurableContext
from aws_durable_execution_sdk_python.execution import (
    InvocationStatus,
    durable_execution,
)
from aws_durable_execution_sdk_python.waits import Duration
from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner

# Deliberately module-global: DurableFunctionTestRunner replays in the SAME
# process, so this flag survives across the suspend/resume boundary — which is
# exactly what makes the divergence reachable. First invocation takes the STEP
# branch; the replay invocation (same process, flag already set) takes the
# CHAINED_INVOKE branch at the same logical position. Against the real service
# the equivalent divergence arises from any replay-visible state change
# (e.g. control flow derived from a corrupted batch rebuild, see JS PR #894).
_STATE = {"diverged_once": False}


@durable_execution
def divergent_handler(_event, context: DurableContext):
    context.step(lambda _: "anchor", name="anchor")
    if not _STATE["diverged_once"]:
        _STATE["diverged_once"] = True
        context.step(lambda _: "first-run-step", name="divergent")
        context.wait(Duration.from_seconds(86400), name="suspend")
    else:
        context.invoke(function_name="some-worker", payload={"x": 1}, name="divergent")
    return {"reached_end": True}


@durable_execution
def stable_handler(_event, context: DurableContext):
    context.step(lambda _: "anchor", name="anchor")
    context.step(lambda _: "stable-step", name="stable")
    context.wait(Duration.from_seconds(86400), name="suspend")
    context.step(lambda _: "post", name="post")
    return {"reached_end": True}


@pytest.mark.example
def test_divergent_replay_fails_closed_with_diagnostic():
    """Type mismatch on replay -> FAILED carrying NonDeterministicExecutionError."""
    _STATE["diverged_once"] = False
    with DurableFunctionTestRunner(handler=divergent_handler) as runner:
        result = runner.run(input='"test"', execution_timeout=30)

    # Never PENDING for a permanent fault (the #867 misclassification).
    assert result.status is not InvocationStatus.PENDING
    assert result.status is InvocationStatus.FAILED, (
        f"Expected FAILED, got {result.status}"
    )
    # Diagnostic must reach the caller with the error type preserved
    # (JS #867 fixed exactly this: bare Error instead of the typed error).
    assert result.error is not None, "FAILED result carries no error object"
    assert "NonDeterministicExecutionError" in (result.error.type or ""), (
        f"Error type not preserved: {result.error.type}"
    )
    assert "type checkpoint='STEP'" in (result.error.message or ""), (
        f"Diagnostic lacks mismatch detail: {result.error.message}"
    )


@pytest.mark.example
def test_stable_replay_control_succeeds():
    """Identical replay across a suspension must succeed (control case)."""
    with DurableFunctionTestRunner(handler=stable_handler) as runner:
        result = runner.run(input='"test"', execution_timeout=30)
    assert result.status is InvocationStatus.SUCCEEDED, (
        f"Control case failed: {result.status}: {result.error}"
    )